#!/usr/bin/env python
# $Id: omd.py 14 2006-08-06 12:08:22Z niels $
#
# Copyright (C) 2006  Niels de Vos
#
# Modified by David Were
#
# This program is free software; you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation; either version 2 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program; if not, write to the Free Software
# Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301, USA
#
"""
Usage of omd: 
"""

from time import sleep
import socket
import logging
from logging.handlers import SysLogHandler

from lib.msg import *
from lib.libomd import *

class OmdMaster(Omd):
    # dict containing key=IP, value=NodeInfo
    __nodeList = nodeList()
    # dict containing key=ip, value=/proc/stat
    #__nodeStatus = dict()
    # create a scheduler - dynamically load random.py?
    scheduler = None
    broadcast = True
    # my IP
    __ip = None
    
    def mainLoop(self, delay):
        """ Scheduler running as a thread.
            Selects a process, selects a node, sends a MigrationRequest.
            The algorithm is implemented in scheduler/*.py:Scheduler
        """
        if not self.scheduler:
            self.log.critical('Programming error: OmdMaster.scheduler not set')
            return
        while self.isRunning:
            try:
                status = StatusMsg()
                self.log.debug('sendMsg: sending message - %s' % status.__class__)
                sendMsg(status)

                proc = self.scheduler.findProcess()
                target = self.scheduler.findNode(self.__nodeList)
                if proc and target:
                    if proc.isMigrated:
                        self.log.debug('Process already migrated: %s %s' % (proc.stat['pid'], proc.stat['comm']))
                        break
                    mr = MigrationRequest(target, proc)
                    sendMsg(mr)
                else:
                    if proc:
                        self.log.debug('No node found for process: %s %s' % (proc.stat['pid'], proc.stat['comm']))
                    elif target:
                        self.log.debug('No process found for node: %s' % (target))
                    else:
                        self.log.debug('No process and no node found')
            except Exception, e:
                self.log.critical('BUG in OmdMaster.mainLoop: %s' % e)
            sleep(delay)
    
    def handleMsg(self, msg):
        """ Handling of the messages, actions to take. """
        if msg.__class__ == InfoRequest:
            # check if it is my own
            if msg.src == msg.dst and not self.__nodeList.contains(msg.src):
                self.__ip = msg.src
                ir = msg.createReply()
                sendMsg(ir)
            else:
                # send InfoReply
                ir = msg.createReply()
                sendMsg(ir)
        elif msg.__class__ == InfoReply:
            # reply on a InfoRequest
            self.__nodeList.addNode(msg.src, msg.cpuinfo, msg.meminfo)
        elif msg.__class__ == StatusMsg:
            # periodic status message, send by all nodes
            # check if we already know this node
            if self.__ip != msg.src and not self.__nodeList.contains(msg.src):
                ir = InfoRequest()
                ir.dst = msg.src
                sendMsg(ir)
            else:
                self.__nodeList.updateNode(msg.src, msg.load, msg.meminfo)
                
        elif msg.__class__ == MigrationRequest:
            # send MigrationReply
            mr = msg.createReply()
            mr.accept = self.acceptor.acceptProcess(msg.proc)
            if mr.accept:
                self.log.debug('accepting %s from %s' % (msg.proc.stat['comm'], msg.src))
            else:
                self.log.debug('not accepting %s from %s' % (msg.proc.stat['comm'], msg.src))
            sendMsg(mr)
        elif msg.__class__ == MigrationReply:
            # reply on a MigrationRequest
            if msg.accept:
                self.log.info('%s accepted %s, %s start migration' % (msg.src, msg.proc.stat['pid'], msg.proc.stat['comm']))
                migrate(msg.proc.stat['pid'], msg.src)
            else:
                self.log.debug('%s did not accept %s' % (msg.src, msg.proc.stat['comm']))
        else:
            # some other message we don't need to handle
            self.log.info('handleMsg: discarding unsupported msg(%s)' % msg.__class__)

def usage():
    print __doc__

if __name__ == '__main__':
    log = logging.getLogger('omd')
    # default algorithm
    schedClass = 'RoundRobin'
    # address to listen on
    listen = ''
    
    try:
        opts, args = getopt.getopt(sys.argv[1:], 'hvl:s', ['help', 'verbose', 'listen=', 'scheduler='])
    except getopt.GetoptError:
        # print help information and exit:
        usage()
        sys.exit(2)

    output = None
    verbose = False
    logLevel = logging.WARNING
    
    logs = list()
    for l in ['scheduler', 'msg', 'omd', 'libomd', 'migration', 'processlist']:
        logs.append(logging.getLogger(l))

    for o, a in opts:
        if o == "-v":
            logLevel -= 10
            if logLevel <= 0:
                logLevel = logging.DEBUG
        elif o in ('-h', '--help'):
            usage()
            sys.exit()
        elif o in ('-l', '--listen'):
            listen = a
        elif o == '-s':
            slh = SysLogHandler()
            for l in logs:
                l.addHandler(slh)
        elif o == '--scheduler':
            schedClass = a

    for l in logs:
        log.setLevel(logLevel)

    omd = OmdMaster()
    try:
        exec('import scheduler.%s' % schedClass)
        exec('omd.scheduler = scheduler.%s.Scheduler()' % schedClass)
        exec('omd.acceptor = scheduler.%s.Acceptor()' % schedClass)
    except Exception, e:
        log.critical('Could not load scheduler %s' % schedClass)
        log.debug(e)
    else:
        omd.scheduler.processList = ProcessList()
        omd.start(listen)

